"""
  此文件用来训练 头车自由驾驶的lstm模型
"""

# !处理路径导入问题（添加绝对路径）！！！
import sys
import os
CODE_INTERNAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) # 生成Code文件夹内部对应的绝对路径
sys.path.append(CODE_INTERNAL_PATH)

# 导入外部包
import numpy as np
from sklearn.model_selection import train_test_split
from keras.models import Sequential
from keras.layers import Dense, LSTM, Dropout, Reshape

# 导入内部包
from utils.read_data import read_extract_follow_data
from utils.lstm import create_lstm_dataset, normalization

EPOCHS = 50 # 迭代次数
BATCH_SIZE = 64 # 一次迭代中的样本数量

FILE_PATH_I80_1_to = "../../Data/Ngsim数据集/I80数据集/3. 提取数据/1. 跟随数据/trajectories-0400-0415_follow.txt"
FILE_PATH_I80_2_to = "../../Data/Ngsim数据集/I80数据集/3. 提取数据/1. 跟随数据/trajectories-0500-0515_follow.txt"

def getData():
  data1 = read_extract_follow_data(FILE_PATH_I80_1_to)
  data2 = read_extract_follow_data(FILE_PATH_I80_2_to)

  data = np.concatenate((data1, data2), axis=0)
  return data

def build_model(input_shape, output_shape):
  model = Sequential([
    LSTM(units=64, input_shape=input_shape, activation="relu", return_sequences=True, dropout=0.2, recurrent_dropout=0.2),
    Dropout(0.3),  # 添加独立Dropout层（训练和预测时均需激活）
    LSTM(units=32, activation="relu", dropout=0.2, recurrent_dropout=0.2),
    Dropout(0.2), # 普通Dropout层（默认仅在训练时激活）
    Dense(output_shape[-1] * output_shape[0], activation="linear"),
    Reshape(output_shape)
  ])
  model.compile(loss="mse", optimizer="adam")
  return model

if __name__ == '__main__':
  # 读取数据
  follow_data = getData()
  print("数据集大小: ", len(follow_data), len(follow_data[0]), len(follow_data[0][0])) # 2727 200 14

  # 归一化
  follow_data, min_max_list = normalization(follow_data, [7, 8, 9, 10, 12], [True, True, False, False, True])
  print("速度的min和max: ", min_max_list[0][0], min_max_list[0][1], min_max_list[1][0], min_max_list[1][1]) # 0.00192 22.38891 0.00055 24.09957
  print("速加度的min和max: ", min_max_list[2][0], min_max_list[2][1], min_max_list[3][0], min_max_list[3][1]) # -3.84815 3.95141 -3.94914 3.86252
  print("相对间距的min和max: ", min_max_list[4][0], min_max_list[4][1]) # 0.013649999999984175 99.53390999999999
  print("数据归一化后: ", follow_data[0][0][7], follow_data[0][0][8], follow_data[0][0][9], follow_data[0][0][10], follow_data[0][0][12]) # 0.168784191175321 0.15845416120655528 -0.3339649928970352 0.18835945240832275 0.13510696214017137

  # 创建训练集和验证集
  n_past, n_future = 30, 10
  X, Y = create_lstm_dataset(follow_data, n_past, n_future, [7, 8, 9, 10, 12], [10])
  x_train, x_val, y_train, y_val = train_test_split(X, Y, test_size=0.3, random_state=42)

  # 构建模型
  input_shape = (n_past, 5)
  output_shape = (n_future, 1)
  model = build_model(input_shape, output_shape)

  # 训练模型
  model.fit(x_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_data=(x_val, y_val), verbose=1)
  model.save("./model/follow_EPOCHS_50_BATCH_SIZE_64_unit_64_32__implement-normalization.keras")